{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training of SST Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torchtext as tt\n",
    "\n",
    "from tqdm import tqdm\n",
    "from pytorch_extras import RAdam, SingleCycleScheduler\n",
    "from pytorch_transformers import GPT2Model, GPT2Tokenizer\n",
    "from deps.torch_train_test_loop.torch_train_test_loop import LoopComponent, TrainTestLoop\n",
    "\n",
    "from models import SSTClassifier"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DEVICE = 'cuda:0'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load pretrained transformer and tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = GPT2Tokenizer.from_pretrained(\n",
    "    'gpt2-large', do_lower_case=False)\n",
    "lang_model = GPT2Model.from_pretrained(\n",
    "    'gpt2-large', output_hidden_states=True, output_attentions=False)\n",
    "lang_model.cuda(device=DEVICE)\n",
    "lang_model.eval()\n",
    "print('Pretrained GPT-2 loaded.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def tokenized_texts_to_embs(tokenized_texts, pad_token=tokenizer.eos_token):\n",
    "    tokenized_texts = [[*tok_seq, pad_token] for tok_seq in tokenized_texts]\n",
    "    lengths = [len(tok_seq) for tok_seq in tokenized_texts]\n",
    "\n",
    "    max_length = max(lengths)\n",
    "    input_toks = [t + [pad_token] * (max_length - l) for t, l in zip(tokenized_texts, lengths)]\n",
    "\n",
    "    input_ids = [tokenizer.convert_tokens_to_ids(tok_seq) for tok_seq in input_toks]\n",
    "    input_ids = torch.tensor(input_ids).to(device=DEVICE)\n",
    "\n",
    "    mask = [[1.0] * length + [0.0] * (max_length - length) for length in lengths]\n",
    "    mask = torch.tensor(mask).to(device=DEVICE)  # [batch sz, num toks]\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        outputs = lang_model(input_ids=input_ids)\n",
    "        embs = torch.stack(outputs[-1], -2)  # [batch sz, n toks, n layers, d emb]\n",
    "\n",
    "    return mask, embs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fine_grained = True  # set to False for binary classification"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SSTFilter():\n",
    "\n",
    "    def __init__(self, remove_neutral=False, remove_dupes=False):\n",
    "        self.remove_neutral, self.remove_dupes  = (remove_neutral, remove_dupes)\n",
    "        self.prev_seen = {}\n",
    "\n",
    "    def __call__(self, sample):\n",
    "        if self.remove_neutral and (sample.label == 'neutral'):\n",
    "            return False\n",
    "        hashable = ''.join(sample.text)\n",
    "        if self.remove_dupes and (hashable in self.prev_seen):\n",
    "            return False\n",
    "        self.prev_seen[hashable] = True\n",
    "        return True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tt.datasets.SST.download(root='.data')  # download if necessary\n",
    "_stoi = { s: i for i, s in enumerate(\n",
    "    ['very negative', 'negative', 'neutral', 'positive', 'very positive'] \\\n",
    "    if fine_grained else ['negative', 'positive']\n",
    ") }\n",
    "TEXT = tt.data.RawField(\n",
    "    preprocessing=tokenizer.tokenize,\n",
    "    postprocessing=tokenized_texts_to_embs,\n",
    "    is_target=False)\n",
    "LABEL = tt.data.RawField(\n",
    "    postprocessing=lambda samples: torch.tensor([_stoi[s] for s in samples], device=DEVICE),\n",
    "    is_target=True)\n",
    "\n",
    "trn_ds = tt.datasets.SST(\n",
    "    '.data/sst/trees/train.txt', TEXT, LABEL, fine_grained=fine_grained, subtrees=True,\n",
    "    filter_pred=SSTFilter(remove_neutral=(not fine_grained), remove_dupes=True))\n",
    "val_ds = tt.datasets.SST(\n",
    "    '.data/sst/trees/dev.txt', TEXT, LABEL, fine_grained=fine_grained, subtrees=False,\n",
    "    filter_pred=SSTFilter(remove_neutral=(not fine_grained), remove_dupes=False))\n",
    "tst_ds = tt.datasets.SST(\n",
    "    '.data/sst/trees/test.txt', TEXT, LABEL, fine_grained=fine_grained, subtrees=False,\n",
    "    filter_pred=SSTFilter(remove_neutral=(not fine_grained), remove_dupes=False))\n",
    "\n",
    "print('Datasets ready.')\n",
    "print('Number of samples: {:,} train phrases, {:,} valid sentences, {:,} test sentences.'\\\n",
    "      .format(len(trn_ds), len(val_ds), len(tst_ds)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training Loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LoopMain(LoopComponent):\n",
    "\n",
    "    def __init__(self, n_classes, device, pct_warmup=0.1, mixup=(0.2, 0.2)):\n",
    "        self.n_classes, self.device, self.pct_warmup = (n_classes, device, pct_warmup)\n",
    "        self.mixup_dist = torch.distributions.Beta(torch.tensor(mixup[0]), torch.tensor(mixup[1]))\n",
    "        self.onehot = torch.eye(self.n_classes, device=self.device)\n",
    "        self.saved_data = []\n",
    "\n",
    "    def on_train_begin(self, loop):\n",
    "        n_iters = len(loop.train_data) * loop.n_epochs\n",
    "        loop.optimizer = RAdam(loop.model.parameters(), lr=5e-4)\n",
    "        loop.scheduler = SingleCycleScheduler(\n",
    "            loop.optimizer, loop.n_optim_steps, frac=self.pct_warmup, min_lr=1e-5)\n",
    "        \n",
    "    def on_grads_reset(self, loop):\n",
    "        loop.model.zero_grad()\n",
    "\n",
    "    def on_forward_pass(self, loop):\n",
    "        model, batch = (loop.model, loop.batch)\n",
    "        mask, embs = batch.text\n",
    "        target_probs = self.onehot[batch.label]\n",
    "\n",
    "        if loop.is_training:\n",
    "            r = self.mixup_dist.sample([len(mask)]).to(device=mask.device)\n",
    "            idx = torch.randperm(len(mask))\n",
    "            mask = mask.lerp(mask[idx], r[:, None])\n",
    "            embs = embs.lerp(embs[idx], r[:, None, None, None])\n",
    "            target_probs = target_probs.lerp(target_probs[idx], r[:, None])\n",
    "\n",
    "        pred_scores, _, _ = model(mask, embs)\n",
    "        _, pred_ids = pred_scores.max(-1)\n",
    "        accuracy = (pred_ids == batch.label).float().mean()\n",
    "\n",
    "        loop.pred_scores, loop.target_probs, loop.accuracy = (pred_scores, target_probs, accuracy)\n",
    "\n",
    "    def on_loss_compute(self, loop):\n",
    "        losses = -loop.target_probs * F.log_softmax(loop.pred_scores, dim=-1)  # CE\n",
    "        loop.loss = losses.sum(dim=-1).mean()  # sum of classes, mean of batch\n",
    "\n",
    "    def on_backward_pass(self, loop):\n",
    "        loop.loss.backward()\n",
    "\n",
    "    def on_optim_step(self, loop):\n",
    "        loop.optimizer.step()\n",
    "        loop.scheduler.step()\n",
    "\n",
    "    def on_batch_end(self, loop):\n",
    "        self.saved_data.append({\n",
    "            'n_samples': len(loop.batch),\n",
    "            'epoch_desc': loop.epoch_desc,\n",
    "            'epoch_num': loop.epoch_num,\n",
    "            'epoch_frac': loop.epoch_num + loop.batch_num / loop.n_batches,\n",
    "            'batch_num' : loop.batch_num,\n",
    "            'accuracy': loop.accuracy.item(),\n",
    "            'loss': loop.loss.item(),\n",
    "            'lr': loop.optimizer.param_groups[0]['lr'],\n",
    "            'momentum': loop.optimizer.param_groups[0]['betas'][0],\n",
    "        })"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LoopProgressBar(LoopComponent):\n",
    "\n",
    "    def __init__(self, item_names=['loss', 'accuracy']):\n",
    "        self.item_names = item_names\n",
    "\n",
    "    def on_epoch_begin(self, loop):\n",
    "        self.total, self.count = ({ name: 0.0 for name in self.item_names }, 0)\n",
    "        self.pbar = tqdm(total=loop.n_batches, desc=f\"{loop.epoch_desc} epoch {loop.epoch_num}\")\n",
    "\n",
    "    def on_batch_end(self, loop):\n",
    "        n = len(loop.batch)\n",
    "        self.count += n\n",
    "        for name in self.item_names:\n",
    "            self.total[name] += getattr(loop, name).item() * n\n",
    "        self.pbar.update(1)\n",
    "        if (not loop.is_training):\n",
    "            means = { f'mean_{name}': self.total[name] / self.count for name in self.item_names }\n",
    "            self.pbar.set_postfix(means)\n",
    "\n",
    "    def on_epoch_end(self, loop):\n",
    "        self.pbar.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initialize and train model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Seed RNG for replicability. Run at least a few times without seeding to measure performance.\n",
    "# torch.manual_seed(<type an int here>)\n",
    "\n",
    "# Make iterators for each split.\n",
    "trn_itr, val_itr, tst_itr = tt.data.Iterator.splits(\n",
    "    (trn_ds, val_ds, tst_ds),\n",
    "    shuffle=True,\n",
    "    batch_size=20,\n",
    "    device=DEVICE)\n",
    "\n",
    "# Initialize model.\n",
    "n_classes = len(_stoi)\n",
    "model = SSTClassifier(\n",
    "    d_depth=lang_model.config.n_layer + 1,\n",
    "    d_emb=lang_model.config.hidden_size,\n",
    "    d_inp=64,\n",
    "    d_cap=2,\n",
    "    n_parts=64,\n",
    "    n_classes=n_classes,\n",
    ")\n",
    "model = model.cuda(device=DEVICE)\n",
    "print('Total number of parameters: {:,}'.format(sum(np.prod(p.shape) for p in model.parameters())))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train model\n",
    "loop = TrainTestLoop(model, [LoopMain(n_classes, DEVICE), LoopProgressBar()], trn_itr, val_itr)\n",
    "loop.train(n_epochs=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loop.test(tst_itr)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
